package db

import (
	"github.com/cheggaaa/pb"
	"github.com/jinzhu/gorm"
	"fmt"
	"cve-db/models"
	log "github.com/sirupsen/logrus"
	"github.com/knqyf263/go-cpe/naming"
	"github.com/knqyf263/go-cpe/common"
	_ "github.com/jinzhu/gorm/dialects/mysql"
)

// 支持的数据库，当前值只适配mysql
const (
	dialectMysql 		= "mysql"
	//dialectSqlite3    = "sqlite3"
	//dialectPostgreSQL = "postgres"
)

// RDB驱动结构体
type RDBDriver struct {
	name string
	conn *gorm.DB
}

// 获取驱动名
func (r *RDBDriver) Name() string {
	return r.name
}

// 返回RDB驱动
func NewRDB(dbType, dsn string, debugSQL bool) (driver *RDBDriver, err error) {
	driver = &RDBDriver{
		name: dbType,
	}

	log.Debugf("Opening DB (%s).", driver.Name())
	if err = driver.OpenDB(dbType, dsn, debugSQL); err != nil {
		return nil, err
	}

	log.Debugf("Migrating DB (%s).", driver.Name())
	if err := driver.MigrateDB(); err != nil {
		return nil, err
	}
	return driver, nil
}

// 新建一个数据库连接
func (r *RDBDriver) OpenDB(dbType, dsn string, debugSQL bool) (err error) {
	r.conn, err = gorm.Open(dbType, dsn)
	if err != nil {
		return fmt.Errorf("Failed to open DB. dbtype: %s, dsn: %s, err: %s",
			dbType, dsn, err)
	}
	if err := r.conn.LogMode(debugSQL).Error; err != nil {
		return err
	}
	r.conn.LogMode(debugSQL)
	return nil
}

// MigrateDB migrates Database
func (r *RDBDriver) MigrateDB() error {
	if err := r.conn.AutoMigrate(
		&models.FeedMeta{},
		&models.CveDetail{},
		&models.NvdJSON{},
		&models.Reference{},
		&models.Cert{},
		&models.Cpe{},
		&models.EnvCpe{},
		&models.Cwe{},
		&models.Affect{},
		&models.Cvss3{},
		&models.Cvss2{},
		&models.Cvss2Extra{},
		&models.Description{},
	).Error; err != nil {
		return fmt.Errorf("Failed to migrate. err: %s", err)
	}

	errs := []error{}

	// CveID
	errs = append(errs, r.conn.Model(&models.CveDetail{}).
		AddIndex("idx_cve_detail_cveid", "cve_id").Error)
	errs = append(errs, r.conn.Model(&models.NvdJSON{}).
		AddIndex("idx_nvd_jsons_cveid", "cve_id").Error)

	// CveDetailID
	errs = append(errs, r.conn.Model(&models.NvdJSON{}).
		AddIndex("idx_nvd_jsons_cve_detail_id", "cve_detail_id").Error)

	// references
	errs = append(errs, r.conn.Model(&models.Reference{}).
		AddIndex("idx_references_nvd_json_id", "nvd_json_id").Error)

	// certs
	errs = append(errs, r.conn.Model(&models.Cert{}).
		AddIndex("idx_certs_nvd_json_id", "nvd_json_id").Error)

	// cpes
	errs = append(errs, r.conn.Model(&models.Cpe{}).
		AddIndex("idx_cpes_nvd_json_id", "nvd_json_id").Error)
	errs = append(errs, r.conn.Model(&models.Cpe{}).
		AddIndex("idx_cpes_uri", "uri").Error)
	errs = append(errs, r.conn.Model(&models.Cpe{}).
		AddIndex("idx_cpes_formatted_string", "formatted_string").Error)
	errs = append(errs, r.conn.Model(&models.Cpe{}).
		AddIndex("idx_cpes_part", "part").Error)
	errs = append(errs, r.conn.Model(&models.Cpe{}).
		AddIndex("idx_cpes_vendor", "vendor").Error)
	errs = append(errs, r.conn.Model(&models.Cpe{}).
		AddIndex("idx_cpes_product", "product").Error)

	// envcpes
	errs = append(errs, r.conn.Model(&models.EnvCpe{}).
		AddIndex("idx_envcpes_cpe_id", "cpe_id").Error)
	errs = append(errs, r.conn.Model(&models.EnvCpe{}).
		AddIndex("idx_envcpes_uri", "uri").Error)
	errs = append(errs, r.conn.Model(&models.EnvCpe{}).
		AddIndex("idx_envcpes_formatted_string", "formatted_string").Error)

	// Cwes
	errs = append(errs, r.conn.Model(&models.Cwe{}).
		AddIndex("idx_cwes_nvd_json_id", "nvd_json_id").Error)

	// Affects
	//TODO add index to vendor, product if needed
	errs = append(errs, r.conn.Model(&models.Affect{}).
		AddIndex("idx_affects_nvd_json_id", "nvd_json_id").Error)

	// Cvss3
	errs = append(errs, r.conn.Model(&models.Cvss3{}).
		AddIndex("idx_cvss3_nvd_json_id", "nvd_json_id").Error)

	// Cvss2
	//errs = append(errs, r.conn.Model(&models.Cvss2{}).
	//	AddIndex("idx_cvsss2_nvd_xml_id", "nvd_xml_id").Error)

	// Cvss2Extra
	errs = append(errs, r.conn.Model(&models.Cvss2Extra{}).
		AddIndex("idx_cvsss2_extra_nvd_json_id", "nvd_json_id").Error)

	// Description
	errs = append(errs, r.conn.Model(&models.Description{}).
		AddIndex("idx_descriptions_nvd_json_id", "nvd_json_id").Error)

	for _, e := range errs {
		if e != nil {
			return fmt.Errorf("Failed to create index. err: %s", e)
		}
	}

	return nil
}

// 从DB中获取CVE信息
func (r *RDBDriver) Get(cveID string) (*models.CveDetail, error) {
	c := models.CveDetail{}
	err := r.conn.Where(&models.CveDetail{CveID: cveID}).First(&c).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("Failed to get cve. err: %s", err)
	}

	if c.ID == 0 {
		return &models.CveDetail{}, nil
	}

	// NVD JSON
	json := models.NvdJSON{}
	err = r.conn.Model(&c).Related(&json, "NvdJSON").Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("Failed to get nvdjson err: %s", err)
	}
	if json.ID == 0 {
		c.NvdJSON = nil
	} else {
		c.NvdJSON = &json
	}
	if json.CveDetailID != 0 && json.ID != 0 {
		descs := []models.Description{}
		err = r.conn.Model(&json).Related(&descs, "Descriptions").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get descriptions err: %s", err)
		}
		c.NvdJSON.Descriptions = descs

		cvss2 := models.Cvss2Extra{}
		err = r.conn.Model(&json).Related(&cvss2, "Cvss2").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get cvss2. err: %s", err)
		}
		c.NvdJSON.Cvss2 = cvss2

		cvss3 := models.Cvss3{}
		err = r.conn.Model(&json).Related(&cvss3, "Cvss3").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get cvss3. err: %s", err)
		}
		c.NvdJSON.Cvss3 = cvss3

		cwes := []models.Cwe{}
		err = r.conn.Model(&json).Related(&cwes, "Cwes").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get cwes. err: %s", err)
		}
		c.NvdJSON.Cwes = cwes

		affects := []models.Affect{}
		err = r.conn.Model(&json).Related(&affects, "Affects").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get affect. err: %s", err)
		}
		c.NvdJSON.Affects = affects

		refs := []models.Reference{}
		err = r.conn.Model(&json).Related(&refs, "References").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get reference. err: %s", err)
		}
		c.NvdJSON.References = refs

		certs := []models.Cert{}
		err = r.conn.Model(&json).Related(&certs, "Certs").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get cert. err: %s", err)
		}
		c.NvdJSON.Certs = certs

		cpes := []models.Cpe{}
		err = r.conn.Model(&json).Related(&cpes, "Cpes").Error
		if err != nil && err != gorm.ErrRecordNotFound {
			return nil, fmt.Errorf("Failed to get cpe. err: %s", err)
		}
		for i, cpe := range cpes {
			envCpes := []models.EnvCpe{}
			err = r.conn.Model(&cpe).Related(&envCpes, "EnvCpes").Error
			if err != nil && err != gorm.ErrRecordNotFound {
				return nil, fmt.Errorf("Failed to get envcpe. err: %s", err)
			}
			cpes[i].EnvCpes = envCpes
		}
		c.NvdJSON.Cpes = cpes
	}

	return &c, nil
}

// 从DB中获取CVEs信息
func (r *RDBDriver) GetMulti(cveIDs []string) (map[string]models.CveDetail, error) {
	cveDetails := map[string]models.CveDetail{}
	for _, cveID := range cveIDs {
		cve, err := r.Get(cveID)
		if err != nil {
			return nil, fmt.Errorf("Failed to get cves. err: %s", err)
		}
		cveDetails[cveID] = *cve
	}
	return cveDetails, nil
}

// 关闭数据库
func (r *RDBDriver) CloseDB() (err error) {
	if err = r.conn.Close(); err != nil {
		log.Errorf("Failed to close DB. Type: %s. err: %s", r.name, err)
		return
	}
	return
}

// 通过 pseudo-CPE 匹配CPEs
func (r *RDBDriver) getMatchingCpes(uri string) ([]models.Cpe, error) {

	// parse wfn, get vendor, product
	specified, err := naming.UnbindURI(uri)
	if err != nil {
		return nil, err
	}
	// select from cpe by vendor, product
	cpes := []models.Cpe{}
	err = r.conn.Where(&models.Cpe{
		CpeBase: models.CpeBase{
			CpeWFN: models.CpeWFN{
				Vendor:  fmt.Sprintf("%s", specified.Get(common.AttributeVendor)),
				Product: fmt.Sprintf("%s", specified.Get(common.AttributeProduct)),
			},
		}}).Find(&cpes).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("Failed to select from cpe by vender,product. err: %s", err)
	}

	log.Debugf("specified: %s", uri)
	filtered := []models.Cpe{}
	checkedIDs := map[uint]bool{}
	for _, cpe := range cpes {
		match, err := match(uri, cpe)
		if err != nil {
			log.Debugf("Failed to compare the version:%s %s cpe_id:%d %#v",
				err, uri, cpe.ID, cpe)

			// Try to exact match by vendor, product and version if the version in CPE is not a semVer style.
			if cpe.NvdJSONID != 0 {
				if _, ok := checkedIDs[cpe.NvdJSONID]; ok {
					continue
				}
				affects := []models.Affect{}
				result := r.conn.Where(&models.Affect{NvdJSONID: cpe.NvdJSONID}).Find(&affects)
				if result.Error != nil && result.Error != gorm.ErrRecordNotFound {
					return nil, fmt.Errorf("Failed to get affects by nvdjson. err: %s", result.Error)
				}

				ok, err := matchExactByAffects(uri, affects)
				if err != nil {
					return nil, fmt.Errorf("Failed to matchExactByAffects. err: %s",err)
				}
				if ok {
					filtered = append(filtered, cpe)
				}
				checkedIDs[cpe.NvdJSONID] = true
			}
		} else if match {
			filtered = append(filtered, cpe)
		}
	}

	return filtered, nil
}

// 通过 pseudo-CPE 获取cve ids
func (r *RDBDriver) GetCveIDsByCpeURI(uri string) ([]string, error) {
	filtered, err := r.getMatchingCpes(uri)
	if err != nil {
		return nil, fmt.Errorf("Failed to get CveIDsByCpeURI. err: %s", err)
	}

	cveIDs := make([]string, len(filtered))

	// The cpes table is de-normalized. So the `First` is correct, don't be mislead.
	for i, cpe := range filtered {
		 if cpe.NvdJSONID != 0 {
			json := models.NvdJSON{}
			err = r.conn.Select("cve_id").Where("ID = ?", cpe.NvdJSONID).First(&json).Error
			if err != nil && err != gorm.ErrRecordNotFound {
				return nil, fmt.Errorf("Failed to get cpes by nvdjsonid. err: %s", err)
			}
			cveIDs[i] = json.CveID
		 }

		// If we didn't find a CVE something is weird.
		if cveIDs[i] == "" {
			log.Infof("Missing cve_id for %s (id: %d)", cpe.URI, cpe.ID)
		}
	}

	return cveIDs, nil
}

// 从DB获取cve信息.
func (r *RDBDriver) GetByCpeURI(uri string) ([]models.CveDetail, error) {
	cveIDs, err := r.GetCveIDsByCpeURI(uri)
	if err != nil {
		return nil, fmt.Errorf("Failed GetByCpeURI. err: %s", err)
	}

	cveDetails, err := r.GetMulti(cveIDs)
	if err != nil {
		return nil, fmt.Errorf("Failed GetMulti. err: %s", err)
	}

	// Convert from map to array
	details := make([]models.CveDetail, len(cveDetails))
	i := 0
	for _, d := range cveDetails {
		details[i] = d
		i++
		log.Debugf("%s", d.CveID)
	}
	return details, nil
}

func rollback(tx *gorm.DB, err error) error {
	if err := tx.Rollback().Error; err != nil {
		return err
	}
	return err
}

// 获取NvdJSON表个数
func (r *RDBDriver) CountNvd() (int, error) {
	var count int
	err := r.conn.Model(&models.NvdJSON{}).Count(&count).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return 0, fmt.Errorf("Failed CountNvd. err: %s", err)
	}
	return count, nil
}

// 插入NvdJson的cve信息到db
func (r *RDBDriver) InsertNvdJSON(cves []models.CveDetail) (err error) {
	log.Infof("Inserting fetched CVEs...")
	bar := pb.New(len(cves))
	bar.Start()

	var refreshedNvds []string
	for chunked := range chunkSlice(cves, 10) {
		tx := r.conn.Begin()
		if tx.Error != nil {
			return tx.Error
		}
		for _, c := range chunked {
			bar.Increment()

			// select old record.
			old := models.CveDetail{}
			result := tx.Where(&models.CveDetail{CveID: c.CveID}).First(&old)
			if result.Error != nil && result.Error != gorm.ErrRecordNotFound {
				return rollback(tx, result.Error)
			}
			if result.RecordNotFound() || old.ID == 0 {
				// log.Debugf("CveDetail is not found: %s", c.CveID)
				if err := tx.Create(&c).Error; err != nil {
					return rollback(tx, fmt.Errorf("Failed to Create CveDetail. err: %s",err))
				}
				refreshedNvds = append(refreshedNvds, c.CveID)
				continue
			}

			if !result.RecordNotFound() {
				// select NvdJSON from db
				json := models.NvdJSON{}
				err = tx.Model(&old).Related(&json, "NvdJSON").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, fmt.Errorf("Failed to select NvdJSON. err: %s", err))
				}

				// 如果NVD JSON还未创建过，插入最新的NVD JSON并和已存在CveDetail表关联
				if json.CveDetailID == 0 {
					log.Debugf("CveDetail found but no NVD JSON: %s", c.CveID)
					c.NvdJSON.CveDetailID = old.ID
					if err := tx.Create(&c.NvdJSON).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to Create NvdJSON. err: %s", err))
					}
					refreshedNvds = append(refreshedNvds, c.CveID)
					continue
				}

				// 如果NVD JSON已经创建过，刷新NVD JSON记录
				// 忽略已经存在的记录并且没有更新过的记录
				if json.LastModifiedDate.Equal(c.NvdJSON.LastModifiedDate) ||
					json.LastModifiedDate.After(c.NvdJSON.LastModifiedDate) {
					log.Debugf("Not modified. old: %s", old.CveID)
					continue
				} else {
					log.Debugf("newer Record found. CveID: %s, old: %s, new: %s",
						c.CveID, json.LastModifiedDate, c.NvdJSON.LastModifiedDate)
				}

				// 删除旧Descriptions
				descs := []models.Description{}
				err = tx.Model(&json).Related(&descs, "Descriptions").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, desc := range descs {
					if err := tx.Unscoped().Delete(desc).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete description. err: %s", err))
					}
				}

				// 删除旧的Cvss2
				cvss2 := []models.Cvss2Extra{}
				err = tx.Model(&json).Related(&cvss2, "Cvss2").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, cvss := range cvss2 {
					if err := tx.Unscoped().Delete(cvss).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete cvss2. err: %s", err))
					}
				}

				// 删除旧的Cvss3
				cvss3 := []models.Cvss3{}
				err = tx.Model(&json).Related(&cvss3, "Cvss3").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, cvss := range cvss3 {
					if err := tx.Unscoped().Delete(cvss).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete cvss3. err: %s", err))
					}
				}

				// 删除旧的CWE
				cwes := []models.Cwe{}
				err = tx.Model(&json).Related(&cwes, "Cwes").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, cwe := range cwes {
					if err := tx.Unscoped().Delete(cwe).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete cwes. err: %s", err))
					}
				}

				// 删除旧的Cpes和EnvEpes
				cpes := []models.Cpe{}
				err = tx.Model(&json).Related(&cpes, "Cpes").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, cpe := range cpes {
					envs := []models.EnvCpe{}
					err = tx.Model(&cpe).Related(&envs, "EnvCpes").Error
					if err != nil && err != gorm.ErrRecordNotFound {
						return rollback(tx, err)
					}
					for _, env := range envs {
						if err := tx.Unscoped().Delete(env).Error; err != nil {
							return rollback(tx, fmt.Errorf("Failed to delete EnvEpes. err: %s", err))
						}
					}

					if err := tx.Unscoped().Delete(cpe).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete cpe. err: %s", err))
					}
				}

				// 删除旧的Affects
				affects := []models.Affect{}
				err = tx.Model(&json).Related(&affects, "Affects").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, r := range affects {
					if err := tx.Unscoped().Delete(r).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete affect. err: %s", err))
					}
				}

				// 删除旧的References
				refs := []models.Reference{}
				err = tx.Model(&json).Related(&refs, "References").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, r := range refs {
					if err := tx.Unscoped().Delete(r).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete references. err: %s", err))
					}
				}

				// 删除旧的Certs
				certs := []models.Cert{}
				err = tx.Model(&json).Related(&certs, "Certs").Error
				if err != nil && err != gorm.ErrRecordNotFound {
					return rollback(tx, err)
				}
				for _, l := range certs {
					if err := tx.Unscoped().Delete(l).Error; err != nil {
						return rollback(tx, fmt.Errorf("Failed to delete certs. err: %s", err))
					}
				}

				// 删除旧的NvdJSON
				if err := tx.Unscoped().Delete(&json).Error; err != nil {
					return rollback(tx, fmt.Errorf("Failed to delete NvdJSON. err: %s", err))
				}

				// 插入NvdJSON
				c.NvdJSON.CveDetailID = old.ID
				if err := tx.Create(&c.NvdJSON).Error; err != nil {
					return rollback(tx, fmt.Errorf("Failed to insert NvdJSON. err: %s", err))
				}
				refreshedNvds = append(refreshedNvds, c.CveID)
			}
		}
		if err := tx.Commit().Error; err != nil {
			return rollback(tx, fmt.Errorf("Failed to commit. err: %s", err))
		}
	}
	bar.Finish()

	log.Infof("Refreshed %d NvdJSONs.", len(refreshedNvds))
	//  log.Debugf("%v", refreshedNvds)
	return nil
}

// 获取feed元数据
func (r *RDBDriver) GetFetchedFeedMeta(url string) (*models.FeedMeta, error) {
	meta := models.FeedMeta{}
	m := &models.FeedMeta{
		URL: url,
	}
	err := r.conn.Where(m).First(&meta).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, err
	}
	return &meta, nil
}

// 插入或更新Feed Hash
func (r *RDBDriver) UpsertFeedHash(mm models.FeedMeta) error {
	meta := models.FeedMeta{}
	m := &models.FeedMeta{
		URL: mm.URL,
	}
	err := r.conn.Where(m).First(&meta).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return err
	}

	tx := r.conn.Begin()
	if tx.Error != nil {
		return tx.Error
	}

	if err == gorm.ErrRecordNotFound {
		m.Hash = mm.Hash
		m.LastModifiedDate = mm.LastModifiedDate
		if err := tx.Create(m).Error; err != nil {
			return rollback(tx, err)
		}
	} else {
		meta.Hash = mm.Hash
		meta.LastModifiedDate = mm.LastModifiedDate
		if err := tx.Save(&meta).Error; err != nil {
			return rollback(tx, err)
		}
	}

	if err := tx.Commit().Error; err != nil {
		return rollback(tx, err)
	}
	return nil
}

// 获取feed元数据列表
func (r *RDBDriver) GetFetchedFeedMetas() ([]models.FeedMeta, error) {
	metas := []models.FeedMeta{}
	err := r.conn.Find(&metas).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, err
	}
	return metas, nil
}
